#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Author  : 1230
# @Email   : xjh_0125@sina.com
# @Time    : 2021/6/21 11:44
# @Software: PyCharm
# @File    : change_bg.py

import colorsys
from PIL import Image
import os
import numpy as np


def replace_color_tran(img, src_clr, dst_clr):
    ''' 通过遍历颜色替换程序
    @param  img:    图像矩阵
    @param  src_clr:    需要替换的颜色(r,g,b)
    @param  dst_clr:    目标颜色        (r,g,b)
    @return             替换后的图像矩阵
    '''
    img_arr = np.asarray(img, dtype=np.double)

    dst_arr = img_arr.copy()
    for i in range(img_arr.shape[1]):
        for j in range(img_arr.shape[0]):
            # tmp = img_arr[j][i]
            # if 10 < tmp[0] < 50 and 100 < tmp[1] < 199 and 170 < tmp[2] < 240:
            # if 225 < tmp[2] < 254:
            # if 30 < tmp[0] < 50 and 140 < tmp[1] < 170:
            if (img_arr[j][i] == src_clr)[0] == True:
                dst_arr[j][i] = dst_clr

    return np.asarray(dst_arr, dtype=np.uint8)


def replace_color(img, src_clr, dst_clr):
    ''' 通过矩阵操作颜色替换程序
    @param  img:    图像矩阵
    @param  src_clr:    需要替换的颜色(r,g,b)
    @param  dst_clr:    目标颜色        (r,g,b)
    @return             替换后的图像矩阵
    '''

    img_arr = np.asarray(img, dtype=np.double)

    # 分离通道
    r_img = img_arr[:, :, 0].copy()
    g_img = img_arr[:, :, 1].copy()
    b_img = img_arr[:, :, 2].copy()

    # 编码
    img = r_img * 256 * 256 + g_img * 256 + b_img
    src_color = src_clr[0] * 256 * 256 + src_clr[1] * 256 + src_clr[2]

    # 索引并替换颜色
    r_img[img == src_color] = dst_clr[0]
    g_img[img == src_color] = dst_clr[1]
    b_img[img == src_color] = dst_clr[2]

    # 合并通道
    dst_img = np.array([r_img, g_img, b_img], dtype=np.uint8)
    # 将数据转换为图像数据(h,w,c)
    dst_img = dst_img.transpose(1, 2, 0)

    return dst_img


class Solution:
    def __init__(self):
        """
        """
        self.filename = 'source.jpg'
        self.out_path = 'target.jpg'

    def change_bg_1(self):
        root = os.getcwd()
        file_path = os.path.join(root, 'img3', self.filename)
        img = Image.open(file_path).convert('RGB')
        img.load()
        target_hue = 0
        r, g, b = img.split()
        result_r, result_g, result_b = [], [], []
        for pixel_r, pixel_g, pixel_b in zip(r.getdata(), g.getdata(), b.getdata()):
            # 转为 HSV 色值
            h, s, v = colorsys.rgb_to_hsv(pixel_r / 255., pixel_b / 255., pixel_g / 255.)

            # 转回 RGB 色系
            rgb = colorsys.hsv_to_rgb(target_hue, s, v)
            pixel_r, pixel_g, pixel_b = [int(x * 255.) for x in rgb]

            # 每个像素点结果保存
            result_r.append(pixel_r)
            result_g.append(pixel_g)
            result_b.append(pixel_b)

        r.putdata(result_r)
        g.putdata(result_g)
        b.putdata(result_b)

        # 合并图片
        image = Image.merge('RGB', (r, g, b))

        # 输出图片
        save_path = os.path.join(root, 'img3', 'target.jpg')
        image.save(save_path)

    def change_bg(self):
        root = os.getcwd()
        file_path = os.path.join(root, 'img3', self.filename)
        out_path = os.path.join(root, 'img3', self.out_path)
        img = Image.open(file_path).convert('RGB')
        dst_img = replace_color(img, (177, 251, 255), (255, 255, 255))
        dst_img = replace_color(dst_img, (234, 252, 238), (255, 255, 255))
        dst_img = replace_color(dst_img, (222, 253, 247), (255, 255, 255))
        dst_img = replace_color(dst_img, (194, 252, 255), (255, 255, 255))
        dst_img = replace_color(dst_img, (175, 239, 251), (255, 255, 255))
        dst_img = replace_color(dst_img, (80, 133, 164), (255, 255, 255))
        dst_img = replace_color(dst_img, (59, 147, 183), (255, 255, 255))
        res_img = Image.fromarray(dst_img)
        res_img.save(out_path)

if __name__ == '__main__':
    s = Solution()
    s.change_bg()
